from omegaconf import DictConfig

from lightning import Trainer
from lightning.pytorch.loggers import WandbLogger

from ecgcmr.signal.sig_datasets.ContrastiveECGLightning import ContrastiveECGDataModule
from ecgcmr.signal.sig_models.ECGViTSimCLR import ECGSimCLR


def train_ecg_contrastive(
        cfg: DictConfig,
        wandb_logger: WandbLogger,
        save_dir: str,
        devices: int = 1
    ):
    datamodule = ContrastiveECGDataModule(cfg=cfg)
    
    model = ECGSimCLR(cfg=cfg, save_dir=save_dir)
    wandb_logger.watch(model, log_graph=False)

    strategy = "ddp" if devices > 1 else "auto"

    trainer = Trainer(
        accelerator="gpu",
        devices=devices,
        strategy=strategy,
        precision="bf16-mixed",
        logger=wandb_logger,
        max_epochs=cfg.max_epochs,
        log_every_n_steps=cfg.log_every_n_steps,
        check_val_every_n_epoch=cfg.check_val_every_n_epoch,
        default_root_dir=save_dir,
        num_sanity_val_steps=0,
        profiler="simple"
    )
    
    trainer.fit(model=model, datamodule=datamodule)